#include <pybind11/pybind11.h>
#include "cutlass_gemm.h"

void gemm_cutlass_wrapper(
    int M, int N, int K,
    uintptr_t A_ptr, int stride_A,
    uintptr_t B_ptr, int stride_B,
    uintptr_t C_ptr, int stride_C,
    uintptr_t D_ptr, int stride_D,
    float alpha, 
    float beta
){
    auto A = reinterpret_cast<void*>(A_ptr);
    auto B = reinterpret_cast<void*>(B_ptr);
    auto C = reinterpret_cast<void*>(C_ptr);
    auto D = reinterpret_cast<void*>(D_ptr);
    gemm_cutlass(
        M, N, K,
        A, stride_A,
        B, stride_B,
        C, stride_C,
        D, stride_D,
        alpha, beta
    );
}

PYBIND11_MODULE(my_cuda_ops, m) {
    m.def("gemm_cutlass", &gemm_cutlass_wrapper, "CUDA operation",
        pybind11::arg("M"), pybind11::arg("N"), pybind11::arg("K"),
        pybind11::arg("A_ptr"), pybind11::arg("stride_A"),
        pybind11::arg("B_ptr"), pybind11::arg("stride_B"),
        pybind11::arg("C_ptr"), pybind11::arg("stride_C"),
        pybind11::arg("D_ptr"), pybind11::arg("stride_D"),
        pybind11::arg("alpha"),
        pybind11::arg("beta")
    );
}
